#!/usr/bin/env python2

import os
import sys
import errno
import getopt

from config import Config
from storage import Storage
from utils import Exp, _dwarn
from diskmap import DiskMap

class MetaNode:
    def __init__(self, storage, config, disk = None):
        self.storage = storage
        self.config = config
        self.test = False

        if (disk == None):
            disk = self.storage.list_node(True);

        self.mmap = DiskMap()
        self.map = DiskMap()
        for (k, v) in disk.items():
            if (v == 'stopped' or 'deleting' in v):
                continue;
            if ("admin" in v or  "meta" in v):
                self.mmap.add(k)
            if ('admin' in v):
                self.admin = k
            self.map.add(k)

        #self.mmap.dump("meta map")
        #self.map.dump("map")

    def __single_node_add(self):
        add = self.map.disk_count() - self.mmap.disk_count()
        add1 = self.config.meta - self.mmap.disk_count()
        if (add == 0 or add1 <= 0): 
            return

        add = add if (add < add1) else add1

        print("skip :" + str(self.mmap.list()))

        lst = self.map.get(add, self.mmap.list(), True)
        for i in lst:
            self.storage.setmeta(i, True)
            self.mmap.add(i)

        print ("new meta disk list (%u) %s" % (self.mmap.disk_count(), str(self.mmap.list())))

    def __single_node_drop(self):
        count = self.mmap.disk_count()
        if (count < self.config.meta):
            return

        lst = self.mmap.list()
        while (len(lst) > self.config.meta):
            disk = lst.pop(0)
            if (disk != self.admin):
                self.storage.dropmeta(disk, True)
                self.mmap.delete(disk)
            else:
                lst.append(disk)

    def __single_rack_add(self):
        add = self.map.node_count() - self.mmap.node_count()
        add1 = self.config.meta - self.mmap.node_count()
        if (add == 0 or add1 <= 0):
            return

        add = add if (add < add1) else add1

        lst = self.map.get(add, self.mmap.list(), False)
        for i in lst:
            self.storage.setmeta(i, True)
            self.mmap.add(i)

        print ("new meta disk list (%u) %s" % (self.mmap.disk_count(), str(self.mmap.list())))

    def __single_rack_drop(self):
        #if (self.mmap.disk_count() + 1 <= self.config.meta):
        #    return

        #print 'self.mmap dump:'
        #self.mmap.dump()

        while (1):
            ent = self.mmap.get_max_node()
            #print 'ent: ', ent
            #print 'ent child: ', ent.child

            if (len(ent.child) == 1):
                break

            while (len(ent.child) > 1):
                disk = ent.child.pop(0)
                
                #print 'disk.id:', disk.id, 'self.admin', self.admin
                if (disk.id != self.admin):
                    try:
                        self.storage.dropmeta(disk.id, True)
                        self.mmap.delete(disk.id)
                    except Exp, e:
                        if (e.errno == errno.EAGAIN):
                            return
                        else:
                            raise
                else:
                    ent.child.append(disk)
        
    def __less_rack_drop(self):
        self.__single_rack_drop()

        if (self.mmap.disk_count() <= self.config.meta):
            return

        disk = self.storage.list_node(True);

        self.mmap = DiskMap()
        for (k, v) in disk.items():
            if (v == 'stopped'):
                continue;

            if ("admin" in v or  "meta" in v):
                self.mmap.add(k)

        while (1):
            ent = self.mmap.get_max_rack()
            if (len(ent.child) <= 1):
                break

            node = ent.child.pop(0)
            disk = node.child[0]
            if (disk.id != self.admin):
                try:
                    self.storage.dropmeta(disk.id, True)
                    self.mmap.delete(disk.id)
                except Exp, e:
                    if (e.errno == errno.EAGAIN):
                        return
                    else:
                        raise
            else:
                ent.child.append(disk)

    def __less_site_add(self):
        self.__more_rack_add()

        add = self.map.node_count() - self.mmap.node_count()
        add1 = self.config.meta - self.mmap.node_count()
        if (add == 0 or add1 <= 0):
            return

        add = add if (add < add1) else add1

        #self.map.dump("dump map ++++++")
        #self.mmap.dump("dump mmap --------")

        left = add
        while (left > 0):
            r = self.mmap.get_min_site().id
            i = self.map.get_insite(r, 1, self.mmap.list())
            #_dmsg("add more rack %s" % (str(i[0])))
            self.mmap.add(i[0])
            #self.mmap.dump("mmap")
            self.storage.setmeta(i[0], True)
            left -= 1

        print ("new meta disk list (%u) %s" % (self.mmap.disk_count(), str(self.mmap.list())))

        while (1):
            ent = self.mmap.get_max_rack()
            if (len(ent.child) <= 1):
                break

            node = ent.child.pop(0)
            disk = node.child[0]
            if (disk.id != self.admin):
                try:
                    self.storage.dropmeta(disk.id, True)
                    self.mmap.delete(disk.id)
                except Exp, e:
                    if (e.errno == errno.EAGAIN):
                        return
                    else:
                        raise
            else:
                ent.child.append(disk)

    def __less_rack_add(self):
        self.__more_rack_add()

        add = self.map.node_count() - self.mmap.node_count()
        add1 = self.config.meta - self.mmap.node_count()
        if (add == 0 or add1 <= 0):
            return

        add = add if (add < add1) else add1

        #self.map.dump("dump map ++++++")
        #self.mmap.dump("dump mmap --------")

        left = add
        skip = []
        while (left > 0):
            while (True):
                try:
                    rack = self.mmap.get_min_rack(skip)
                    if (rack is None):
                        _dwarn("rack is None")
                        return

                    r = rack.id
                    i = self.map.get_inrack(r, 1, self.mmap.list())
                    break
                except Exp, e:
                    if (e.errno == errno.ENOSPC):
                        skip.append(r)
                        continue
                    raise

            #_dmsg("add more rack %s" % (str(i[0])))
            self.mmap.add(i[0])
            #self.mmap.dump("mmap")
            self.storage.setmeta(i[0], True)
            left -= 1

        print ("new meta disk list (%u) %s" % (self.mmap.disk_count(), str(self.mmap.list())))

    def __less_rack_drop(self):
        self.__single_rack_drop()

        if (self.mmap.disk_count() <= self.config.meta):
            return

        disk = self.storage.list_node(True);

        self.mmap = DiskMap()
        for (k, v) in disk.items():
            if (v == 'stopped'):
                continue;

            if ("admin" in v or  "meta" in v):
                self.mmap.add(k)

        while (1):
            ent = self.mmap.get_max_rack()
            if (len(ent.child) <= 1):
                break

            node = ent.child.pop(0)
            disk = node.child[0]
            if (disk.id != self.admin):
                try:
                    self.storage.dropmeta(disk.id, True)
                    self.mmap.delete(disk.id)
                except Exp, e:
                    if (e.errno == errno.EAGAIN):
                        return
                    else:
                        raise
            else:
                ent.child.append(disk)


    def __more_rack_add(self):
        add = self.map.rack_count() - self.mmap.rack_count()
        add1 = self.config.meta - self.mmap.rack_count()
        if (add == 0 or add1 <= 0):
            return

        add = add if (add < add1) else add1

        diff = self.map.rack_count() - self.mmap.rack_count()
        diff = diff if (diff < add) else add

        #_dwarn("new rack %u %u" % (add, diff))

        if (diff > 0):
            lst = self.map.get(diff, self.mmap.list())
            #_dmsg("add new rack %s" % (str(lst)))
            for i in lst:
                self.storage.setmeta(i, True)
                self.mmap.add(i)

            add -= diff

        print ("new meta disk list (%u) %s" % (self.mmap.disk_count(), str(self.mmap.list())))

    def __more_rack_drop(self):
        count = self.mmap.disk_count()
        if (count < self.config.meta):
            return

        lst = self.mmap.list()
        while (len(lst) > self.config.meta):
            disk = lst.pop(0)
            if (disk != self.admin):
                self.storage.dropmeta(disk, True)
                self.mmap.delete(disk)
            else:
                lst.append(disk)

    def __single_node(self):
        self.__single_node_add()
        self.__single_node_drop()

    def __single_rack(self):
        self.__single_rack_add()
        self.__single_rack_drop()

    def __more_rack(self):
        self.__more_rack_add()
        self.__more_rack_drop()

    def __less_rack(self):
        self.__less_rack_add()
        self.__less_rack_drop()

    def __more_site(self):
        self.__more_rack_add()
        self.__more_rack_drop()

    def __less_site(self):
        self.__less_site_add()
        self.__less_rack_drop()

    def __clean_up(self):
        if (self.test):
            return

        _metas = self.mmap.list()
        if(len(_metas) == 0):
            return

        stat = self.storage.nodeinfo(_metas[0])
        metas = stat['metas'].split(',')
        if (metas[-1] == ''):
            metas.pop(-1)

        for i in metas:
            if (i not in (_metas)):
                _dwarn("clean up meta %s\n" % (i))
                try:
                    self.storage.dropmeta(i)
                except Exp, e:
                    return

    def balance(self):
        self.__clean_up()

        if (self.map.node_count() == 1): #only one node exist
            self.__single_node()
        elif (self.map.rack_count() == 1): #only one rack exist
            self.__single_rack()
        elif (self.map.site_count() == 1):
            if (self.map.rack_count() < self.config.meta): #rack count less then config
                self.__less_rack()
            else:
                self.__more_rack() #rack count more then config
        else:
            if (self.map.site_count() < self.config.meta): #rack count less then config
                self.__less_site()
            else:
                self.__more_site() #rack count more then config

class FakeStorage():
    def setmeta(self, disk, xx):
        print("-----set %s meta" % (disk))

    def dropmeta(self, disk, xx):
        print("-----drop %s meta" % (disk))


def __test_single_node1():
    config = Config()
    storage = FakeStorage()

    disk = {
        "d1.r1.h1/0" : "admin",
        "d1.r1.h1/1" : "normal",
        "d1.r1.h1/2" : "normal",
        "d1.r1.h1/3" : "normal",
        "d1.r1.h1/4" : "normal",
        }

    move = MetaNode(storage, config, disk)
    move.test = True
    move.balance()

def __test_single_node2():
    config = Config()
    storage = FakeStorage()

    disk = {
        "d1.r1.h1/0" : "admin",
        "d1.r1.h1/1" : "normal",
        "d1.r1.h1/2" : "normal",
        "d1.r1.h1/3" : "normal",
        "d1.r1.h1/4" : "normal",
        "d1.r1.h1/5" : "normal",
        "d1.r1.h1/6" : "normal",
        "d1.r1.h1/7" : "normal",
        "d1.r1.h1/8" : "normal",
        "d1.r1.h1/9" : "normal",
        "d1.r1.h1/10" : "normal",
        "d1.r1.h1/11" : "normal",
        }

    move = MetaNode(storage, config, disk)
    move.test = True
    move.balance()
    
    
def __test_single_rack_2node():
    config = Config()
    storage = FakeStorage()

    disk = {
        "d1.r1.h1/0" : "normal",
        "d1.r1.h1/1" : "normal",
        "d1.r1.h1/2" : "admin",
        "d1.r1.h1/3" : "normal",
        "d1.r1.h1/4" : "normal",
        "d1.r1.h2/5" : "normal",
        "d1.r1.h2/6" : "normal",
        "d1.r1.h2/7" : "normal",
        "d1.r1.h2/8" : "meta",
        "d1.r1.h2/9" : "normal",
        "d1.r1.h2/10" : "normal",
        "d1.r1.h2/11" : "normal",
        }

    move = MetaNode(storage, config, disk)
    move.test = True
    move.balance()
    
def __test_single_rack_3node():
    config = Config()
    storage = FakeStorage()

    disk = {
        "d1.r1.h1/0" : "normal",
        "d1.r1.h1/1" : "normal",
        "d1.r1.h1/2" : "normal",
        "d1.r1.h1/3" : "normal",
        "d1.r1.h1/4" : "normal",
        "d1.r1.h2/5" : "normal",
        "d1.r1.h2/6" : "normal",
        "d1.r1.h3/7" : "normal",
        "d1.r1.h3/8" : "normal",
        "d1.r1.h3/9" : "normal",
        "d1.r1.h3/10" : "normal",
        "d1.r1.h3/11" : "normal",
        }

    move = MetaNode(storage, config, disk)
    move.test = True
    move.balance()

def __test_single_rack_3node_2disk():
    config = Config()
    storage = FakeStorage()

    disk = {
            'site1.rack1.node185/0': 'meta',
            'site1.rack1.node185/1': 'normal',

            'site1.rack1.node186/0': 'normal',
            'site1.rack1.node186/1': 'normal',

            'site1.rack1.node184/0': 'admin',
            'site1.rack1.node184/1': 'normal',
        }

    move = MetaNode(storage, config, disk)
    move.test = True
    move.balance()
 
    
def __test_single_rack_4node():
    config = Config()
    storage = FakeStorage()

    disk = {
        "d1.r1.h1/0" : "normal",
        "d1.r1.h1/1" : "normal",
        "d1.r1.h1/2" : "normal",
        "d1.r1.h1/3" : "normal",
        "d1.r1.h1/4" : "normal",
        "d1.r1.h2/5" : "normal",
        "d1.r1.h2/6" : "normal",
        "d1.r1.h3/7" : "normal",
        "d1.r1.h3/8" : "normal",
        "d1.r1.h4/9" : "normal",
        "d1.r1.h4/10" : "normal",
        "d1.r1.h4/11" : "normal",
        }

    move = MetaNode(storage, config, disk)
    move.test = True
    move.balance()
    
def __test_single_rack_5node():
    config = Config()
    storage = FakeStorage()

    disk = {
        "d1.r1.h1/0" : "normal",
        "d1.r1.h1/1" : "normal",
        "d1.r1.h1/2" : "normal",
        "d1.r1.h1/3" : "normal",
        "d1.r1.h1/4" : "normal",
        "d1.r1.h2/5" : "normal",
        "d1.r1.h2/6" : "normal",
        "d1.r1.h3/7" : "normal",
        "d1.r1.h3/8" : "normal",
        "d1.r1.h4/9" : "normal",
        "d1.r1.h4/10" : "normal",
        "d1.r1.h5/11" : "normal",
        }

    move = MetaNode(storage, config, disk)
    move.test = True
    move.balance()
    
def __test_2rack():
    config = Config()
    storage = FakeStorage()

    disk = {
        "d1.r1.h1/0" : "normal",
        "d1.r1.h1/1" : "normal",
        "d1.r1.h2/0" : "normal",
        "d1.r1.h2/1" : "normal",
        "d1.r1.h3/0" : "normal",
        "d1.r1.h3/1" : "normal",

        "d1.r2.h1/0" : "normal",
        "d1.r2.h1/1" : "normal",
        "d1.r2.h2/0" : "normal",
        "d1.r2.h2/1" : "normal",
        "d1.r2.h3/0" : "normal",
        "d1.r2.h3/1" : "normal",
        }

    print ("test list:" + str(disk))
    move = MetaNode(storage, config, disk)
    move.test = True
    move.balance()
    
def __test_3rack():
    config = Config()
    storage = FakeStorage()

    disk = {
        "d1.r1.h1/0" : "normal",
        "d1.r1.h1/1" : "normal",
        "d1.r1.h2/0" : "normal",
        "d1.r1.h2/1" : "normal",
        "d1.r1.h3/0" : "normal",
        "d1.r1.h3/1" : "normal",

        "d1.r2.h1/0" : "normal",
        "d1.r2.h1/1" : "normal",
        "d1.r2.h2/0" : "normal",
        "d1.r2.h2/1" : "normal",
        "d1.r2.h3/0" : "normal",
        "d1.r2.h3/1" : "normal",

        "d1.r3.h1/0" : "normal",
        "d1.r3.h1/1" : "normal",
        "d1.r3.h2/0" : "normal",
        "d1.r3.h2/1" : "normal",
        "d1.r3.h3/0" : "normal",
        "d1.r3.h3/1" : "normal",
        }

    print ("test list:" + str(disk))
    move = MetaNode(storage, config, disk)
    move.test = True
    move.balance()
    
def __test_9rack():
    config = Config()
    storage = FakeStorage()

    disk = {
        "d1.r1.h1/0" : "admin",
        "d1.r1.h1/1" : "normal",
        "d1.r1.h2/0" : "normal",
        "d1.r1.h2/1" : "normal",

        "d1.r2.h1/0" : "normal",
        "d1.r2.h1/1" : "normal",
        "d1.r2.h2/0" : "meta",
        "d1.r2.h2/1" : "normal",

        "d1.r3.h1/0" : "normal",
        "d1.r3.h1/1" : "normal",
        "d1.r3.h2/0" : "normal",
        "d1.r3.h2/1" : "normal",

        "d1.r4.h1/0" : "normal",
        "d1.r4.h1/1" : "normal",
        "d1.r4.h2/0" : "normal",
        "d1.r4.h2/1" : "normal",

        "d1.r5.h1/0" : "normal",
        "d1.r5.h1/1" : "normal",
        "d1.r5.h2/0" : "normal",
        "d1.r5.h2/1" : "normal",

        "d1.r6.h1/0" : "normal",
        "d1.r6.h1/1" : "normal",
        "d1.r6.h2/0" : "normal",
        "d1.r6.h2/1" : "normal",

        "d1.r7.h1/0" : "normal",
        "d1.r7.h1/1" : "normal",
        "d1.r7.h2/0" : "normal",
        "d1.r7.h2/1" : "normal",

        "d1.r8.h1/0" : "normal",
        "d1.r8.h1/1" : "normal",
        "d1.r8.h2/0" : "normal",
        "d1.r8.h2/1" : "normal",

        "d1.r9.h1/0" : "normal",
        "d1.r9.h1/1" : "normal",
        "d1.r9.h2/0" : "normal",
        "d1.r9.h2/1" : "normal",
        }

    print ("test list:" + str(disk))
    move = MetaNode(storage, config, disk)
    move.test = True
    move.balance()
    
def __test_2site_2rack_3node():
    config = Config()
    storage = FakeStorage()

    disk = {
        "d1.r1.h1/0" : "normal",
        "d1.r1.h1/1" : "normal",
        "d1.r1.h1/3" : "normal",

        "d1.r1.h2/0" : "normal",
        "d1.r1.h2/1" : "normal",
        "d1.r1.h2/2" : "normal",

        "d2.r2.h3/0" : "admin",
        "d2.r2.h3/1" : "normal",
        "d2.r2.h3/2" : "normal",
        }

    print ("test list:" + str(disk))
    move = MetaNode(storage, config, disk)
    move.test = True
    move.balance()

def __test_1site_3rack_7node():
    config = Config()
    storage = FakeStorage()

    disk = {
            'datacenter1.rack1.node11/0': 'admin',
            'datacenter1.rack1.node11/1': 'normal',

            'datacenter1.rack1.node12/0': 'normal',
            'datacenter1.rack1.node12/1': 'normal',
            'datacenter1.rack1.node12/2': 'normal',

            'datacenter1.rack1.node16/0': 'normal',
            'datacenter1.rack1.node16/1': 'normal',
            'datacenter1.rack1.node16/2': 'normal',
            'datacenter1.rack1.node16/3': 'normal',

            'datacenter1.rack2.node14/0': 'normal',
            'datacenter1.rack2.node14/1': 'meta',
            'datacenter1.rack2.node14/2': 'normal',

            'datacenter1.rack2.node13/0': 'normal',
            'datacenter1.rack2.node13/1': 'normal',
            'datacenter1.rack2.node13/2': 'normal',

            'datacenter1.rack2.node17/0': 'normal',
            'datacenter1.rack2.node17/1': 'normal',
            'datacenter1.rack2.node17/2': 'normal',
            'datacenter1.rack2.node17/3': 'normal',

            'datacenter1.rack3.node19/0': 'meta',
        }

    print ("test list:" + str(disk))
    move = MetaNode(storage, config, disk)
    move.test = True
    move.balance()

def __test_2site_2rack_4node():
    config = Config()
    storage = FakeStorage()

    disk = {
        "d1.r1.h1/0" : "normal",
        "d1.r1.h1/1" : "normal",
        "d1.r1.h1/3" : "normal",

        "d1.r1.h2/0" : "normal",
        "d1.r1.h2/1" : "normal",
        "d1.r1.h2/2" : "normal",

        "d2.r2.h3/0" : "admin",
        "d2.r2.h3/1" : "normal",
        "d2.r2.h3/2" : "normal",

        "d2.r2.h4/0" : "normal",
        "d2.r2.h4/1" : "normal",
        "d2.r2.h4/2" : "normal",
        }

    print ("test list:" + str(disk))
    move = MetaNode(storage, config, disk)
    move.test = True
    move.balance()

def __test_1site_3rack_7node():
    config = Config()
    storage = FakeStorage()

    disk = {
            'datacenter1.rack1.node11/0': 'normal',
            'datacenter1.rack1.node11/1': 'normal',

            'datacenter1.rack1.node12/0': 'normal',
            'datacenter1.rack1.node12/1': 'admin',
            'datacenter1.rack1.node12/2': 'normal',

            'datacenter1.rack1.node16/0': 'meta',
            'datacenter1.rack1.node16/1': 'normal',
            'datacenter1.rack1.node16/2': 'normal',
            'datacenter1.rack1.node16/3': 'normal',

            'datacenter1.rack2.node13/0': 'meta',
            'datacenter1.rack2.node13/1': 'normal',
            'datacenter1.rack2.node13/2': 'normal',

            'datacenter1.rack2.node14/0': 'normal',
            'datacenter1.rack2.node14/1': 'normal',
            'datacenter1.rack2.node14/2': 'normal',

            'datacenter1.rack2.node17/0': 'normal',
            'datacenter1.rack2.node17/1': 'meta',
            'datacenter1.rack2.node17/2': 'normal',
            'datacenter1.rack2.node17/3': 'normal',

            'datacenter1.rack3.node15/0': 'normal',
            'datacenter1.rack3.node15/1': 'meta',

            'datacenter1.rack3.node19/0': 'meta',
        }

    print ("test list:" + str(disk))
    move = MetaNode(storage, config, disk)
    move.test = True
    move.balance()



def __test():
    __test_single_node1()
    __test_single_node2()

    __test_single_rack_2node()
    __test_single_rack_3node()
    __test_single_rack_4node()
    __test_single_rack_5node()
    __test_2rack()
    __test_3rack()
    __test_9rack()
    __test_2site_2rack_3node()
    __test_2site_2rack_4node()
    
    __test_1site_3rack_7node()
    __test_1site_3rack_7node()

    print '*'*100
    __test_single_rack_3node_2disk()

def __balance():
    config = Config()
    storage = Storage(config)
    move = MetaNode(storage, config, None)
    move.balance()

def usage():
    print ("usage:")
    print (sys.argv[0] + " --balance")
    print (sys.argv[0] + " --test")

def main():
    try:
        opts, args = getopt.getopt(
                sys.argv[1:], 
                'hv', ['test', 'balance', 'help']
                )
    except getopt.GetoptError, err:
        print str(err)
        usage()
        exit(errno.EINVAL)

    for o, a in opts:
        if o in ('--help'):
            usage()
            exit(0)
        elif o == '--test':
            __test()
        elif o == '--balance':
            __balance()
        else:
            usage()
            exit(0)

if __name__ == '__main__':
    if (len(sys.argv) == 1):
        usage()
    else:
        main()
